import paddle
from paddle.vision.models import resnet50
from paddle.vision.transforms import Compose, Resize, ToTensor, Normalize
from PIL import Image
import numpy as np
import paddle.nn.functional as F
import json

class InsectPredictor:
    def __init__(self, model_path, class_index_path, threshold=0.5):
        self.threshold = threshold
        # 加载类别索引
        with open(class_index_path, 'r') as f:
            class_to_idx = json.load(f)
        self.idx_to_class = {idx: label for label, idx in class_to_idx.items()}

        self.model = resnet50(pretrained=False, num_classes=len(self.idx_to_class))
        self.model.set_state_dict(paddle.load(model_path))
        self.model.eval()
        self.transform = Compose([
            Resize((224, 224)),
            ToTensor(),
            Normalize(mean=[0.485, 0.456, 0.406],
                      std=[0.229, 0.224, 0.225])
        ])

    def predict(self, image: Image.Image):
        img = self.transform(np.array(image)).unsqueeze(0)
        logits = self.model(img)
        probs = F.softmax(logits, axis=1).numpy()[0]
        pred = probs.argmax()
        confidence = probs[pred]
        if confidence < self.threshold:
            return "不确定类别（置信度低）", confidence
        else:
            return self.idx_to_class[pred], confidence
